#!/usr/bin/env python3

import rospy
import time
import math
import numpy as np

from std_msgs.msg import Header
from geometry_msgs.msg import PoseStamped, TwistStamped
from mavros_msgs.msg import State, PositionTarget
from mavros_msgs.srv import CommandBool, SetMode, CommandBoolResponse, SetModeResponse


class MockPX4:
    def __init__(self):
        rospy.init_node('mock_px4_node', anonymous=True)
        rospy.loginfo("Mock PX4 Node Initializing...")

        self.rate = rospy.Rate(50)  # Simulate at a higher frequency for smoother motion

        # --- Simulated Drone State ---
        self.current_pose = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])  # x, y, z, roll, pitch, yaw
        self.current_velocity = np.array([0.0, 0.0, 0.0])  # vx, vy, vz
        self.current_yaw_rate = 0.0

        self.px4_state = State()
        self.px4_state.armed = False
        self.px4_state.connected = True
        self.px4_state.mode = "GUIDED"  # Start in GUIDED mode

        # --- MAVROS Service Server ---
        self.arming_srv = rospy.Service('/uav1/mavros/cmd/arming', CommandBool, self._arming_cb)
        self.set_mode_srv = rospy.Service('/uav1/mavros/set_mode', SetMode, self._set_mode_cb)

        # --- MAVROS Publishers ---
        self.pose_pub = rospy.Publisher('/uav1/mavros/local_position/pose', PoseStamped, queue_size=1)
        self.state_pub = rospy.Publisher('/uav1/mavros/state', State, queue_size=1)

        # --- MAVROS Subscribers ---
        rospy.Subscriber('/uav1/mavros/setpoint_raw/local', PositionTarget, self._setpoint_raw_cb)

        rospy.loginfo("Mock PX4 Node Ready.")

    def _arming_cb(self, req: CommandBool):
        if self.px4_state.mode == "OFFBOARD":  # Only allow arming in OFFBOARD in simulation for safety
            self.px4_state.armed = req.value
            rospy.loginfo(f"Mock PX4: Arming set to {self.px4_state.armed}")
            return CommandBoolResponse(success=True)
        else:
            rospy.logwarn("Mock PX4: Cannot arm in non-OFFBOARD mode.")
            return CommandBoolResponse(success=False)

    def _set_mode_cb(self, req: SetMode):
        if req.custom_mode == "OFFBOARD":
            self.px4_state.mode = "OFFBOARD"
            rospy.loginfo("Mock PX4: Mode set to OFFBOARD")
            return SetModeResponse(mode_sent=True)
        else:
            rospy.logwarn(f"Mock PX4: Only OFFBOARD mode change is simulated. Requested: {req.custom_mode}")
            return SetModeResponse(mode_sent=False)

    def _setpoint_raw_cb(self, msg: PositionTarget):
        # This is a simplified interpretation of PositionTarget.type_mask
        # In real MAVROS, it's more complex, but for simulation, we'll extract relevant parts.

        # Position Control (if Px, Py, Pz are not ignored)
        if not (msg.type_mask & PositionTarget.IGNORE_PX):
            self.current_velocity[0] = (msg.position.x - self.current_pose[
                0]) * self.rate.h * 0.5  # P-controller for position
        if not (msg.type_mask & PositionTarget.IGNORE_PY):
            self.current_velocity[1] = (msg.position.y - self.current_pose[1]) * self.rate.h * 0.5
        if not (msg.type_mask & PositionTarget.IGNORE_PZ):
            self.current_velocity[2] = (msg.position.z - self.current_pose[2]) * self.rate.h * 0.5

        # Velocity Control (if Vx, Vy, Vz are not ignored)
        if not (msg.type_mask & PositionTarget.IGNORE_VX):
            self.current_velocity[0] = msg.velocity.x
        if not (msg.type_mask & PositionTarget.IGNORE_VY):
            self.current_velocity[1] = msg.velocity.y
        if not (msg.type_mask & PositionTarget.IGNORE_VZ):
            self.current_velocity[2] = msg.velocity.z

        # Yaw Rate Control (if Yaw is ignored, Yaw Rate is used)
        if (msg.type_mask & PositionTarget.IGNORE_YAW) and (msg.yaw_rate != 0):
            self.current_yaw_rate = msg.yaw_rate
        else:
            self.current_yaw_rate = 0.0

        # Apply simple velocity limits
        max_speed = 5.0  # m/s
        self.current_velocity = np.clip(self.current_velocity, -max_speed, max_speed)
        self.current_yaw_rate = np.clip(self.current_yaw_rate, -math.pi, math.pi)  # 1 rad/s

    def _update_drone_pose(self):
        dt = 1.0 / self.rate.h

        # Update position based on velocity (in body frame, then rotate to NED)
        # Assuming current_velocity is in drone's body frame for simplicity of IBVS interpretation
        # In PX4, setpoint_raw/local is typically in NED. Let's assume NED for now.
        self.current_pose[0] += self.current_velocity[0] * dt  # X (North)
        self.current_pose[1] += self.current_velocity[1] * dt  # Y (East)
        self.current_pose[2] += self.current_velocity[2] * dt  # Z (Down)

        # Update yaw based on yaw rate
        self.current_pose[5] += self.current_yaw_rate * dt  # Yaw (Z-axis rotation)
        self.current_pose[5] = math.fmod(self.current_pose[5], 2 * math.pi)  # Keep yaw within 0 to 2pi

        # Publish updated pose
        pose_msg = PoseStamped()
        pose_msg.header.stamp = rospy.Time.now()
        pose_msg.header.frame_id = "map"  # Or "odom", "local_ned"
        pose_msg.pose.position.x = self.current_pose[0]
        pose_msg.pose.position.y = self.current_pose[1]
        pose_msg.pose.position.z = self.current_pose[2]  # Z-down

        # Convert Euler to Quaternion for pose.orientation
        q = tf.quaternion_from_euler(self.current_pose[3], self.current_pose[4], self.current_pose[5])
        pose_msg.pose.orientation.x = q[0]
        pose_msg.pose.orientation.y = q[1]
        pose_msg.pose.orientation.z = q[2]
        pose_msg.pose.orientation.w = q[3]
        self.pose_pub.publish(pose_msg)

        # Publish updated PX4 state
        self.px4_state.header.stamp = rospy.Time.now()
        self.state_pub.publish(self.px4_state)

    def run(self):
        rospy.loginfo("Mock PX4 Node Running.")
        while not rospy.is_shutdown():
            self._update_drone_pose()
            self.rate.sleep()


if __name__ == '__main__':
    try:
        mock_px4 = MockPX4()
        mock_px4.run()
    except rospy.ROSInterruptException:
        pass
